'''
Data structure for performing intersect queries on a set of intervals which
preserves all information about the intervals (unlike bitset projection methods).

:Authors: James Taylor (james@jamestaylor.org),
          Ian Schenk (ian.schenck@gmail.com),
          Brent Pedersen (bpederse@gmail.com)
'''

# Historical note:
#    This module original contained an implementation based on sorted endpoints
#    and a binary search, using an idea from Scott Schwartz and Piotr Berman.
#    Later an interval tree implementation was implemented by Ian for Galaxy's
#    join tool (see `bx.intervals.operations.quicksect.py`). This was then
#    converted to Cython by Brent, who also added support for
#    upstream/downstream/neighbor queries. This was modified by James to
#    handle half-open intervals strictly, to maintain sort order, and to
#    implement the same interface as the original Intersecter.

#cython: cdivision=True

from    libc.stdlib   cimport rand, RAND_MAX
from    libc.math     cimport ceil, log
from    libc.limits   cimport INT_MIN, INT_MAX

import operator


cdef inline int imax2(int a, int b):
  if b > a:
    return b
  return a

cdef inline int imax3(int a, int b, int c):
  if b > a:
    if c > b:
      return c
    return b
  if a > c:
    return a
  return c

cdef inline int imin3(int a, int b, int c):
  if b < a:
    if c < b:
      return c
    return b
  if a < c:
    return a
  return c

cdef inline int imin2(int a, int b):
  if b < a:
    return b
  return a

cdef float nlog = -1.0/log(0.5)

cdef class IntervalNode(object):
  '''
  A single node of an `IntervalTree`.

  NOTE: Unless you really know what you are doing, you probably should us
        `IntervalTree` rather than using this directly.
  '''
  cdef float         priority
  cdef public object value
  cdef public int    start, end
  cdef int           minend, maxend, minstart
  cdef IntervalNode  cleft, cright, croot

  property left:
    def __get__(self):
      return <object>self.cleft if self.cleft is not EmptyNode else None

  property right:
    def __get__(self):
      return <object>self.cright if self.cright is not EmptyNode else None

  property root:
    def __get__(self):
      return <object>self.croot if self.croot is not EmptyNode else None

  def __repr__(self):
    if self.value is not None:
      return 'IntervalNode(%i, %i, %r)' % (self.start, self.end, self.value)
    else:
      return 'IntervalNode(%i, %i)' % (self.start, self.end)

  def __cinit__(IntervalNode self, int start, int end, object value=None):
    # Python lacks the binomial distribution, so we convert a
    # uniform into a binomial because it naturally scales with
    # tree size.  Also, python's uniform is perfect since the
    # upper limit is not inclusive, which gives us undefined here.
    self.priority   = ceil(nlog * log(-1.0/(1.0 * rand()/RAND_MAX - 1)))
    self.start      = start
    self.end        = end
    self.value      = value
    self.maxend     = end
    self.minstart   = start
    self.minend     = end
    self.cleft      = EmptyNode
    self.cright     = EmptyNode
    self.croot      = EmptyNode

  cpdef IntervalNode insert(IntervalNode self, int start, int end, object value=None):
    '''
    Insert a new IntervalNode into the tree of which this node is
    currently the root. The return value is the new root of the tree (which
    may or may not be this node!)
    '''
    cdef IntervalNode croot = self

    if start > self.start or (start==self.start and end > self.end):
      # insert to cright tree
      if self.cright is not EmptyNode:
        self.cright = self.cright.insert(start, end, value)
      else:
        self.cright = IntervalNode(start, end, value)
      # rebalance tree
      if self.priority < self.cright.priority:
        croot = self.rotate_left()
    else:
      # insert to cleft tree
      if self.cleft is not EmptyNode:
        self.cleft = self.cleft.insert(start, end, value)
      else:
        self.cleft = IntervalNode(start, end, value)
        # rebalance tree
        if self.priority < self.cleft.priority:
          croot = self.rotate_right()

    croot.set_ends()
    self.cleft.croot  = croot
    self.cright.croot = croot
    return croot

  cdef IntervalNode rotate_right(IntervalNode self):
    cdef IntervalNode croot = self.cleft
    self.cleft   = self.cleft.cright
    croot.cright = self
    self.set_ends()
    return croot

  cdef IntervalNode rotate_left(IntervalNode self):
    cdef IntervalNode croot = self.cright
    self.cright = self.cright.cleft
    croot.cleft = self
    self.set_ends()
    return croot

  cdef inline void set_ends(IntervalNode self):
    if self.cright is not EmptyNode and self.cleft is not EmptyNode:
      self.maxend   = imax3(self.end,   self.cright.maxend,   self.cleft.maxend  )
      self.minend   = imin3(self.end,   self.cright.minend,   self.cleft.minend  )
      self.minstart = imin3(self.start, self.cright.minstart, self.cleft.minstart)
    elif self.cright is not EmptyNode:
      self.maxend   = imax2(self.end,   self.cright.maxend  )
      self.minend   = imin2(self.end,   self.cright.minend  )
      self.minstart = imin2(self.start, self.cright.minstart)
    elif self.cleft is not EmptyNode:
      self.maxend   = imax2(self.end,   self.cleft.maxend  )
      self.minend   = imin2(self.end,   self.cleft.minend  )
      self.minstart = imin2(self.start, self.cleft.minstart)

  def intersect(self, int start, int end):
    '''
    given a start and a end, return a list of features
    falling within that range
    '''
    cdef list results = []
    self._intersect(start, end, results)
    return results

  cdef void _intersect(IntervalNode self, int start, int end, list results):
    # Left subtree
    if self.cleft is not EmptyNode and self.cleft.maxend > start:
        self.cleft._intersect(start, end, results)
    # This interval
    if self.end > start and self.start < end:
        results.append(self)
    # Right subtree
    if self.cright is not EmptyNode and self.start < end:
        self.cright._intersect(start, end, results)

  cdef void _seek_left(IntervalNode self, int position, list results, int n, int max_dist):
    # we know we can bail in these 2 cases.
    if max_dist and self.maxend + max_dist < position:
      return
    if self.minstart > position:
      return

    # the ordering of these 3 blocks makes it so the results are
    # ordered nearest to farest from the query position
    if self.cright is not EmptyNode:
      self.cright._seek_left(position, results, n, max_dist)

    d = position - self.end

    if d>=0 and (max_dist==0 or d < max_dist):
      results.append(self)

    # TODO: can these conditionals be more stringent?
    if self.cleft is not EmptyNode:
      self.cleft._seek_left(position, results, n, max_dist)

  cdef void _seek_right(IntervalNode self, int position, list results, int n, int max_dist):
    # we know we can bail in these 2 cases.
    if self.maxend < position:
      return
    if max_dist and self.minstart - max_dist > position:
      return

    # the ordering of these 3 blocks makes it so the results are
    # ordered nearest to farest from the query position
    if self.cleft is not EmptyNode:
      self.cleft._seek_right(position, results, n, max_dist)

    d = self.start - position

    if d>=0 and (max_dist==0 or d<max_dist):
      results.append(self)

    if self.cright is not EmptyNode:
      self.cright._seek_right(position, results, n, max_dist)

  cpdef min(self):
    '''
    find smallest interval coordinate
    '''
    return self.minstart

  cpdef max(self):
    '''
    find largest interval coordinate
    '''
    return self.maxend

  cpdef before(self, position, int n=0, int max_dist=0, sort=True):
    '''
    find n features with a start > than `position`
    f: a Interval object (or anything with an `end` attribute)
    n: the number of features to return
    max_dist: the maximum distance to look before giving up.
    '''
    cdef list results = []
    # use start - 1 becuase .before() assumes strictly left-of
    self._seek_left(position - 1, results, n, max_dist)
    if n and len(results) == n:
      return results

    r = results

    if n or sort:
      r.sort(key=operator.attrgetter('end','start'), reverse=True)
      if n:
        r=r[:n]

    return r

  cpdef after(self, position, int n=0, int max_dist=0, sort=True):
    '''
    find n features with a end < than position
    f: a Interval object (or anything with a `start` attribute)
    n: the number of features to return
    max_dist: the maximum distance to look before giving up.
    '''
    cdef list results = []
    # use end + 1 becuase .after() assumes strictly right-of
    self._seek_right(position + 1, results, n, max_dist)
    if n and len(results) == n:
      return results

    r = results

    if n or sort:
      r.sort(key=operator.attrgetter('start','end'))
      if n:
        r=r[:n]

    return r

  def traverse(self, func):
      self._traverse(func)

  cdef void _traverse(IntervalNode self, object func):
      if self.cleft is not EmptyNode: self.cleft._traverse(func)
      func(self)
      if self.cright is not EmptyNode: self.cright._traverse(func)


cdef IntervalNode EmptyNode = IntervalNode(0, 0)


cdef class IntervalTree:
  '''
  Data structure for performing window intersect queries on a set of
  of possibly overlapping 1d intervals.

  Usage
  =====

  Create an empty IntervalTree

  >>> from intervaltree import IntervalTree
  >>> itree = IntervalTree()

  An interval is a start and end position and a value (possibly None).
  You can add any object as an interval:

  >>> itree.insert(0, 10, 'food')
  >>> itree.insert(3, 7, dict(foo='bar'))

  >>> itree.find(2, 5)
  [IntervalNode(0, 10, 'food'), IntervalNode(3, 7, {'foo': 'bar'})]

  If the object has start and end attributes (like the Interval class) there
  is are some shortcuts:

  >>> itree = IntervalTree()
  >>> itree.insert(0, 10)
  >>> itree.insert(3, 7)
  >>> itree.insert(3, 40)
  >>> itree.insert(13, 50)

  >>> itree.find(30, 50)
  [IntervalNode(3, 40), IntervalNode(13, 50)]
  >>> itree.find(100, 200)
  []
  '''

  cdef IntervalNode root

  def __cinit__(self):
    root = None

  def __iter__(self):
    return iter(self.find_values())

  # ---- Position based interfaces -----------------------------------------

  def insert(self, int start, int end, object value=None):
    '''
    Insert the interval [start,end) associated with value `value`.
    '''
    if self.root:
      self.root = self.root.insert(start, end, value)
    else:
      self.root = IntervalNode(start, end, value)

  def find(self, start=INT_MIN, end=INT_MAX):
    '''
    Return a sorted list of all intervals overlapping [start,end).
    '''
    if self.root is None:
      return []
    return self.root.intersect(start, end)

  def find_values(self, start=INT_MIN, end=INT_MAX):
    '''
    Return a sorted list of all intervals overlapping [start,end).
    '''
    if self.root is None:
      return []
    return [ r.value for r in self.root.intersect(start, end) ]

  def before(self, position, num_intervals=0, max_dist=0):
    '''
    Find `num_intervals` intervals that lie before `position` and are no
    further than `max_dist` positions away
    '''
    if self.root is None:
      return []
    return self.root.before(position, num_intervals, max_dist)

  def after(self, position, num_intervals=0, max_dist=0):
    '''
    Find `num_intervals` intervals that lie after `position` and are no
    further than `max_dist` positions away
    '''
    if self.root is None:
      return []
    return self.root.after(position, num_intervals, max_dist)
